import heapq

from sortedcontainers import SortedList


class DinnerPlates:
    def __init__(self, capacity: int):
        self.capacity = capacity
        self.stacks = [[]]  # 栈列表
        self.empty = [0]  # 所有有空位的栈下标
        self.values = SortedList()  # 所有包含值的栈下标

    def push(self, val: int) -> None:
        idx = self.empty[0]
        self.stacks[idx].append(val)
        if idx not in self.values:
            self.values.add(idx)

        if len(self.stacks[idx]) == self.capacity:
            heapq.heappop(self.empty)

            if not self.empty:
                heapq.heappush(self.empty, len(self.stacks))
                self.stacks.append([])

    def pop(self) -> int:
        if len(self.values) == 0:
            return -1

        idx = self.values[-1]

        if len(self.stacks[idx]) == self.capacity:
            heapq.heappush(self.empty, idx)

        ans = self.stacks[idx].pop()
        if len(self.stacks[idx]) == 0:
            self.values.remove(idx)
        return ans

    def popAtStack(self, idx: int) -> int:
        if len(self.stacks) <= idx:
            return -1
        if len(self.stacks[idx]) == 0:
            return -1
        if len(self.stacks[idx]) == self.capacity:
            heapq.heappush(self.empty, idx)
        ans = self.stacks[idx].pop()
        if len(self.stacks[idx]) == 0:
            self.values.remove(idx)
        return ans


if __name__ == "__main__":
    obj = DinnerPlates(2)
    obj.push(1)
    obj.push(2)
    obj.push(3)
    obj.push(4)
    obj.push(5)
    print(obj.popAtStack(0))  # 2
    obj.push(20)
    obj.push(21)
    print(obj.popAtStack(0))  # 20
    print(obj.popAtStack(2))  # 21
    print(obj.pop())  # 5
    print(obj.pop())  # 4
    print(obj.pop())  # 3
    print(obj.pop())  # 1
    print(obj.pop())  # -1
